(*
  Copyright (c) 2009 Barry Schwartz

  Permission is hereby granted, free of charge, to any person
  obtaining a copy of this software and associated documentation
  files (the "Software"), to deal in the Software without
  restriction, including without limitation the rights to use,
  copy, modify, merge, publish, distribute, sublicense, and/or sell
  copies of the Software, and to permit persons to whom the
  Software is furnished to do so, subject to the following
  conditions:

  The above copyright notice and this permission notice shall be
  included in all copies or substantial portions of the Software.

  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
  OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
  HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
  WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
  OTHER DEALINGS IN THE SOFTWARE.
*)

open XNum
  
type compare_result = | Eq | Lt | Gt

type 'a compare = 'a -> 'a -> compare_result

type 'a lin_form =
  { const : num; terms : (num * 'a) list; compare : 'a compare
  }

let of_num compare n = { const = n; terms = []; compare = compare; }
  
let lin_zero compare = of_num compare num_zero
  
let of_scaled_unknown compare a x =
  if eq_num a num_zero
  then lin_zero compare
  else { const = num_zero; terms = [ (a, x) ]; compare = compare; }
  
let of_unknown compare x = of_scaled_unknown compare num_one x
  
let of_terms compare xs =
  {
    const = num_zero;
    terms = List.filter (fun (a, _) -> not (eq_num a num_zero)) xs;
    compare = compare;
  }
  
let is_constant lin = match lin.terms with | [] -> true | _ -> false
  
let coefficient lin x =
  let rec iter terms =
    match terms with
    | [] -> num_zero
    | (a, y) :: ys ->
        (match lin.compare x y with
         | Lt -> num_zero
         | Eq -> a
         | Gt -> iter ys)
  in iter lin.terms
  
let remove_first_term lin =
  match lin.terms with
  | [] -> invalid_arg "LinForm.remove_first_term"
  | _ :: xs -> { (lin) with terms = xs; }
  
let add l0 l1 =
  let rec add_summands xs ys =
    match (xs, ys) with
    | ([], _) -> ys
    | (_, []) -> xs
    | ((c, x) :: xx, (d, y) :: yy) ->
        (match l0.compare x y with
         | Lt -> (c, x) :: (add_summands xx ys)
         | Gt -> (d, y) :: (add_summands xs yy)
         | Eq ->
             let e = add_num c d
             in
               if eq_num e num_zero
               then add_summands xx yy
               else (e, x) :: (add_summands xx yy))
  in
    {
      const = add_num l0.const l1.const;
      terms = add_summands l0.terms l1.terms;
      compare = l0.compare;
    }
  
let add_const lin c = { (lin) with const = add_num lin.const c; }
  
let add_unknown lin a x = add lin (of_scaled_unknown lin.compare a x)
  
let scale c lin =
  if eq_num c num_zero
  then lin_zero lin.compare
  else
    {
      const = mult_num c lin.const;
      terms = List.map (fun (b, x) -> ((mult_num c b), x)) lin.terms;
      compare = lin.compare;
    }
  
let lin_comb a l0 b l1 =
  if eq_num a num_zero
  then scale b l1
  else
    if eq_num b num_zero
    then scale a l0
    else
      (let rec add_summands xs ys =
         match (xs, ys) with
         | ([], _) -> ys
         | (_, []) -> xs
         | ((c, x) :: xx, (d, y) :: yy) ->
             (match l0.compare x y with
              | Lt -> ((mult_num a c), x) :: (add_summands xx ys)
              | Gt -> ((mult_num b d), y) :: (add_summands xs yy)
              | Eq ->
                  let e = add_num (mult_num a c) (mult_num b d)
                  in
                    if eq_num e num_zero
                    then add_summands xx yy
                    else (e, x) :: (add_summands xx yy))
       in
         {
           const = add_num (mult_num a l0.const) (mult_num b l1.const);
           terms = add_summands l0.terms l1.terms;
           compare = l0.compare;
         })
  
let sub l0 l1 = lin_comb num_one l0 num_minus_one l1
  
let map comp f lin =
  {
    const = lin.const;
    terms = List.map (fun (b, x) -> (b, (f x))) lin.terms;
    compare = comp;
  }
  

